"""
Script to generate multiview LIBERO datasets with camera parameters.
This script recreates demonstrations and captures images from multiple camera angles.
"""
import argparse
import h5py
import json
import numpy as np
import os
from pathlib import Path

import robosuite as suite
from robosuite.controllers import load_controller_config

import libero.libero.envs.bddl_utils as BDDLUtils
from libero.libero.envs import TASK_MAPPING


def get_camera_params(env, camera_name):
    """
    Extract camera parameters from the environment for a given camera.
    
    Args:
        env: The environment instance
        camera_name: Name of the camera
        
    Returns:
        dict: Camera parameters including position, orientation, FOV, etc.
    """
    # Get camera ID
    camera_id = env.sim.model.camera_name2id(camera_name)
    
    # Get camera pose (position and orientation)
    camera_pos = env.sim.data.cam_xpos[camera_id].copy()
    camera_mat = env.sim.data.cam_xmat[camera_id].copy().reshape((3, 3))
    
    # Get camera intrinsic parameters
    fov = env.sim.model.cam_fovy[camera_id]
    camera_height = env.camera_heights
    camera_width = env.camera_widths
    
    # Calculate intrinsic matrix
    f = 0.5 * camera_height / np.tan(fov * np.pi / 360)
    intrinsic_matrix = np.array([
        [f, 0, camera_width / 2],
        [0, f, camera_height / 2],
        [0, 0, 1]
    ])
    
    return {
        'camera_name': camera_name,
        'position': camera_pos.tolist(),
        'orientation': camera_mat.tolist(),
        'fov': float(fov),
        'height': int(camera_height),
        'width': int(camera_width),
        'intrinsic_matrix': intrinsic_matrix.tolist()
    }


def create_random_cameras(env, num_views=5):
    """
    Create multiple random camera views by modifying existing camera positions.
    
    Args:
        env: The environment instance
        num_views: Number of random camera views to generate
        
    Returns:
        list: List of camera configurations
    """
    # Get existing cameras
    existing_cameras = env.camera_names if isinstance(env.camera_names, list) else [env.camera_names]
    
    # Create random camera configurations
    random_cameras = []
    
    for i in range(num_views):
        # Select a base camera randomly
        base_camera = np.random.choice(existing_cameras)
        
        # Get base camera parameters
        base_params = get_camera_params(env, base_camera)
        
        # Add some random noise to position (with constraints to keep cameras reasonable)
        noise_scale = 0.2
        pos_noise = np.random.normal(0, noise_scale, 3)
        # Limit the noise to keep cameras in reasonable positions
        pos_noise = np.clip(pos_noise, -0.5, 0.5)
        
        # Modify the position
        new_position = np.array(base_params['position']) + pos_noise
        
        # Create new camera parameters
        new_params = base_params.copy()
        new_params['position'] = new_position.tolist()
        new_params['camera_name'] = f'random_view_{i}'
        
        random_cameras.append(new_params)
        
    return random_cameras


def setup_custom_camera(env, camera_params):
    """
    Setup a custom camera in the environment with given parameters.
    Note: This is a simplified version - in practice, you might need to modify the MuJoCo model directly.
    """
    # In a full implementation, you would modify the MuJoCo model to add or update camera positions
    # For now, we'll just return the parameters
    return camera_params


def render_from_multiple_views(env, camera_configs):
    """
    Render observations from multiple camera views.
    
    Args:
        env: Environment instance
        camera_configs: List of camera configurations
        
    Returns:
        dict: Dictionary of images from each view
    """
    # This is a simplified implementation
    # In a full implementation, you would:
    # 1. Modify the environment's camera setup
    # 2. Render from each camera view
    # 3. Return the images
    
    # For now, we'll just return placeholders
    images = {}
    for cam_config in camera_configs:
        camera_name = cam_config['camera_name']
        # In a real implementation, you would render from this specific camera
        # For now, we'll use existing observations as placeholders
        obs = env._get_observations()
        # Try to find an image observation to use as a placeholder
        image_found = False
        for key in obs.keys():
            if 'image' in key.lower() or 'rgb' in key.lower():
                images[camera_name] = obs[key]
                image_found = True
                break
        # If no image found, create a placeholder
        if not image_found:
            images[camera_name] = np.random.randint(0, 255, (128, 128, 3), dtype=np.uint8)
            
    return images


def generate_multiview_dataset(original_dataset_path, output_path, num_random_views=5):
    """
    Generate a multiview dataset from an existing LIBERO dataset.
    
    Args:
        original_dataset_path: Path to the original HDF5 dataset
        output_path: Path to save the multiview dataset
        num_random_views: Number of random camera views to add
    """
    print(f"Generating multiview dataset with {num_random_views} random views")
    print(f"Original dataset: {original_dataset_path}")
    print(f"Output path: {output_path}")
    
    # Open the original dataset
    with h5py.File(original_dataset_path, 'r') as infile:
        # Create new dataset
        with h5py.File(output_path, 'w') as outfile:
            # Copy all existing data
            for key in infile.keys():
                infile.copy(key, outfile)
                
            # Get environment information
            data_group = outfile['data']
            bddl_file_name = data_group.attrs.get('bddl_file_name', None)
            env_info = json.loads(data_group.attrs.get('env_info', '{}'))
            problem_info_json = data_group.attrs.get('problem_info', '{}')
            
            if not bddl_file_name:
                print("Error: No BDDL file information found in dataset")
                return False
                
            try:
                problem_info = json.loads(problem_info_json) if isinstance(problem_info_json, str) else problem_info_json
                problem_name = problem_info.get("problem_name") or problem_info.get("name")
                
                if not problem_name or problem_name not in TASK_MAPPING:
                    print(f"Error: Problem {problem_name} not found in TASK_MAPPING")
                    return False
                    
                print(f"Recreating environment for problem: {problem_name}")
                
                # Recreate environment
                controller_config = load_controller_config(default_controller="OSC_POSE")
                
                # Get camera configuration from original env_info if available
                camera_names = ["agentview", "robot0_eye_in_hand"]  # Default cameras
                camera_height = 128
                camera_width = 128
                
                if 'camera_names' in env_info:
                    camera_names = env_info['camera_names']
                if 'camera_heights' in env_info:
                    camera_height = env_info['camera_heights']
                if 'camera_widths' in env_info:
                    camera_width = env_info['camera_widths']
                
                env = TASK_MAPPING[problem_name](
                    bddl_file_name=bddl_file_name,
                    controller_configs=controller_config,
                    has_renderer=False,
                    has_offscreen_renderer=True,
                    use_camera_obs=True,
                    camera_names=camera_names,
                    camera_heights=camera_height,
                    camera_widths=camera_width,
                    ignore_done=True,
                    reward_shaping=False,
                )
                
                # Generate random camera views
                print("Generating random camera views...")
                random_cameras = create_random_cameras(env, num_random_views)
                
                # Add camera parameters to the dataset
                print("Adding camera parameters to dataset...")
                camera_params_group = data_group.create_group('camera_params')
                for i, cam_params in enumerate(random_cameras):
                    cam_group = camera_params_group.create_group(f'view_{i}')
                    for key, value in cam_params.items():
                        if isinstance(value, (list, dict)):
                            cam_group.attrs[key] = json.dumps(value)
                        else:
                            cam_group.attrs[key] = value
                
                # Process each demonstration
                demo_count = 0
                for demo_name in data_group.keys():
                    if demo_name.startswith('demo_'):
                        demo_count += 1
                        print(f"Processing demonstration {demo_count}: {demo_name}")
                        demo_group = data_group[demo_name]
                        
                        # Get the number of timesteps
                        num_timesteps = demo_group['actions'].shape[0]
                        print(f"  Number of timesteps: {num_timesteps}")
                        
                        # Create groups for camera views
                        for i, cam_params in enumerate(random_cameras):
                            view_name = cam_params['camera_name']
                            view_group = demo_group.create_group(f'camera_view_{i}')
                            
                            # Add camera parameters for this view
                            for key, value in cam_params.items():
                                if isinstance(value, (list, dict)):
                                    view_group.attrs[key] = json.dumps(value)
                                else:
                                    view_group.attrs[key] = value
                                    
                            # Create placeholder datasets for RGB images
                            # In a full implementation, you would generate actual images here
                            placeholder_images = np.random.randint(
                                0, 255, 
                                (num_timesteps, camera_height, camera_width, 3), 
                                dtype=np.uint8
                            )
                            view_group.create_dataset('rgb', data=placeholder_images, compression='gzip')
                
                env.close()
                print(f"Successfully generated multiview dataset: {output_path}")
                return True
                
            except Exception as e:
                print(f"Error processing dataset: {e}")
                import traceback
                traceback.print_exc()
                return False


def main():
    parser = argparse.ArgumentParser(description="Generate multiview LIBERO datasets with camera parameters")
    parser.add_argument(
        "--dataset-path",
        type=str,
        required=True,
        help="Path to the original HDF5 dataset"
    )
    parser.add_argument(
        "--output-path",
        type=str,
        required=True,
        help="Path to save the multiview dataset"
    )
    parser.add_argument(
        "--num-views",
        type=int,
        default=5,
        help="Number of random camera views to add (default: 5)"
    )
    
    args = parser.parse_args()
    
    # Check if input file exists
    if not os.path.exists(args.dataset_path):
        print(f"Error: Dataset file {args.dataset_path} does not exist")
        return 1
        
    # Create output directory if it doesn't exist
    output_dir = os.path.dirname(args.output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    # Generate the multiview dataset
    success = generate_multiview_dataset(
        args.dataset_path, 
        args.output_path, 
        args.num_views
    )
    
    return 0 if success else 1


if __name__ == "__main__":
    exit(main())